from keras.datasets import mnist
from keras.utils import np_utils

(X_train, y_train), (X_test, y_test) = mnist.load_data()

num_train_samples = X_train.shape[0]
num_test_samples = y_test.shape[0]
image_height = X_train.shape[1]
image_width = X_train.shape[2]
print("X_train.shape: " + str(X_train.shape))
print("X_test.shape: " + str(X_test.shape))
print("y_train.shape: " + str(y_train.shape))
print("y_test.shape: " + str(y_test.shape))
print("y_train sample 5 value: " + str(y_train[5]))
print("Train samples: " + str(num_train_samples))
print("Test samples: " + str(num_test_samples))
print("Image height: " + str(image_height))
print("Image width: " + str(image_width))

# Conv2D expects 4 dimmensions; the last one is the number of channels
num_channels = 1

X_train = X_train.reshape(num_train_samples,image_height,
                          image_width,num_channels)
X_test = X_test.reshape(num_test_samples,image_height,
                        image_width,num_channels)
  
X_train = X_train / 255   # values [0..1] improve results
X_test = X_test / 255
    
y_train = np_utils.to_categorical(y_train)
y_test = np_utils.to_categorical(y_test)

print("categorical y_train shape: " + str(y_train.shape))
print("categorical y_train sample 5 value: " + str(y_train[5]))
num_classes = y_test.shape[1]
print("Number of classes: " + str(num_classes))
